import os
import re
import cv2
import time
import json

def generate_video_with_speed_opencv(
        input_dir: str,
        output_file: str,
        base_framerate: int = 30,
):
    """
    使用OpenCV生成变速视频，并记录每个变速区间的开始时间点
    :param input_dir: 图片目录（包含类似0.jpg的图片和range.json配置）
    :param output_file: 输出视频路径
    :param base_framerate: 基准帧率（输出视频帧率）
    """
    # 读取速度配置文件
    json_path = os.path.join(input_dir, "range.json")
    if not os.path.exists(json_path):
        raise ValueError(f"未找到速度配置文件: {json_path}")
    
    with open(json_path, "r") as f:
        try:
            speed_data = json.load(f)
        except Exception as e:
            raise ValueError(f"解析JSON失败: {e}")

    # 转换JSON格式为速度区间元组
    speed_segments = []
    for seg in speed_data:
        if "start" not in seg or "factor" not in seg:
            raise ValueError("range.json格式错误: 必须包含start和factor字段")
        
        start = int(seg["start"])
        end = seg.get("end")
        factor = float(seg["factor"])
        speed_segments.append((
            start,
            int(end) if (end is not None and end != "null") else None,
            factor
        ))

    # 预处理速度区间（按起始帧排序）
    speed_segments.sort(key=lambda x: x[0])

    # 初始化记录每个区间开始时间的数据结构
    seg_time_info = [{'triggered': False, 'start_time': None} for _ in speed_segments]

    # 获取所有图片并排序
    frame_files = []
    for f in os.listdir(input_dir):
        if f == "range.json":
            continue  # 跳过配置文件
        
        match = re.match(r"(\d+)\.jpg$", f)
        if match:
            frame_num = int(match.group(1))
            frame_files.append((frame_num, os.path.join(input_dir, f)))

    if not frame_files:
        raise ValueError("未找到任何图片文件")

    # 按帧号排序
    frame_files.sort(key=lambda x: x[0])

    # 生成帧序列
    output_frames = []
    for frame_num, file_path in frame_files:
        current_factor = 1.0
        for seg in speed_segments:
            start, end, factor = seg
            if end is None:
                if frame_num >= start:
                    current_factor = factor
                    break
            else:
                if start <= frame_num <= end:
                    current_factor = factor
                    break

        # 检查并记录变速区间的开始时间
        for i, seg in enumerate(speed_segments):
            if seg_time_info[i]['triggered']:
                continue
            start, end, factor = seg
            in_segment = False
            if end is None:
                in_segment = frame_num >= start
            else:
                in_segment = start <= frame_num <= end
            if in_segment:
                current_output_frame_count = len(output_frames)
                start_time = current_output_frame_count / base_framerate
                seg_time_info[i]['start_time'] = start_time
                seg_time_info[i]['triggered'] = True

        # 读取图片并计算重复次数
        img = cv2.imread(file_path)
        if img is None:
            raise ValueError(f"无法读取图片: {file_path}")

        repeat = max(1, int(round(1 / current_factor))) if current_factor > 0 else 1
        output_frames.extend([img] * repeat)

    # 构建时间记录结果
    time_result = []
    for i, seg in enumerate(speed_segments):
        start, end, factor = seg
        time_entry = {
            "start": start,
            "end": end,
            "factor": factor,
            "start_time": seg_time_info[i]['start_time']
        }
        if(factor<1.0):
            time_result.append(time_entry)

    # 保存时间记录到JSON文件
    json_output_path = os.path.splitext(output_file)[0] + "_time.json"
    with open(json_output_path, 'w') as f:
        json.dump(time_result, f, indent=2, ensure_ascii=False)
    print(f"变速时间点已保存至: {json_output_path}")

    # 写入视频
    if not output_frames:
        raise ValueError("没有可输出的帧")

    height, width, _ = output_frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    
    out = cv2.VideoWriter(output_file, fourcc, base_framerate, (width, height))
    for frame in output_frames:
        out.write(frame)
    out.release()
    print(f"视频生成成功: {output_file}")


# 使用示例
if __name__ == "__main__":
    start = time.time()
    generate_video_with_speed_opencv(
        input_dir="/data/control_communication/results/video",
        output_file="resultnew.mp4",
        base_framerate=30
    )
    end = time.time()
    print(f"耗时: {end-start:.2f}秒")